#include <iostream>
#include <fstream>
#include <Eigen/Core>
#include <Eigen/Array>
#include <vector>
#include <math.h>

USING_PART_OF_NAMESPACE_EIGEN

using namespace std;


void activationFunction(const MatrixXd& x, MatrixXd & result);

// Returns a random positive integer in [0,lim-1]
int randInt(int lim);

void readImage();
int loadinit=1;
int imgSize = 512;
int numImages = 10;
string filename = "olsh.dat";
char* parafile;
char fstr[20];
vector<MatrixXd> images;
int patchSize=8;
int inputSize=patchSize*patchSize;
int hiddenSize=30;
int outputSize=inputSize;
int it=100000;//1000000000;
int printintv=100;
int outputnum=10;
double alpha=0.01;
double beta=5;
double ro=-0.96;
double lambda=0.002;
double rhoDecay=0.99;
MatrixXf W1(hiddenSize,inputSize+1);
MatrixXf W2(outputSize,hiddenSize+1);
VectorXf a1(inputSize+1);
VectorXf a2(hiddenSize+1);
VectorXf a3(outputSize);
VectorXf rho(hiddenSize);
VectorXf delta3(outputSize);
VectorXf delta2(hiddenSize);

void readImage()
{
	ifstream indata;
	double num; // variable for input value
	
	indata.open(filename.c_str());
	if (!indata) {
		cerr << "Error: file could not be opened" << endl;
		return;
	}
	
	for (int i = 0; i < numImages; ++i) {
		//will m be destroyed when the function returns?
		MatrixXd m(imgSize,imgSize);
		for (int r = 0; r < imgSize; ++r) {
			for (int c = 0; c < imgSize; ++c) {
				if (indata.eof()) {
					cerr << "Error: ran out of input values on (" << r << "," << c << ")" << endl;
					return;
				}
				indata >> num;
				m(r,c) = num;
			}
		}
		images.push_back(m);
	}
	indata.close();
}


void normalizeInput()
{
	double maxin=-100;
	double minin=100;
	for (int i = 0; i < numImages; ++i) {
		for (int r = 0; r < imgSize; ++r) {
			for (int c = 0; c < imgSize; ++c) {
				double num=images[i](r,c);
				if (maxin<num)
					maxin=num;
				if (minin>num)
					minin=num;
			}
		}
	}
	//normalize the input to [-1,1]
	for (int mi=0;mi<numImages;mi++) {
		for (int r=0;r<imgSize;r++) {
			for (int c=0;c<imgSize;c++) {
				images[mi](r,c)=(images[mi](r,c)-(maxin+minin)/2.0)/((maxin-minin)/2.0);
			}
		}
	}
}

void trimInput()
{
	//normalize the input to [-1,1]
	for (int mi=0;mi<numImages;mi++) {
		for (int r=0;r<imgSize;r++) {
			for (int c=0;c<imgSize;c++) {
				double num=images[mi](r,c);
				if(num>1)
					images[mi](r,c)=1;
				if(num<-1)
					images[mi](r,c)=-1;
			}
		}
	}
}

void tanhNormalize()
{
	//normalize the input to [-1,1]
	for (int mi=0;mi<numImages;mi++) {
		for (int r=0;r<imgSize;r++) {
			for (int c=0;c<imgSize;c++) {
				double num=images[mi](r,c);
				images[mi](r,c)=tanh(num);
			}
		}
	}
}

void zeroMeanScaleInput()
{
	double maxin=-100;
	double minin=100;
	for (int i = 0; i < numImages; ++i) {
		for (int r = 0; r < imgSize; ++r) {
			for (int c = 0; c < imgSize; ++c) {
				double num=images[i](r,c);
				if (maxin<num)
					maxin=num;
				if (minin>num)
					minin=num;
			}
		}
	}
	//normalize the input to [-1,1]
	for (int mi=0;mi<numImages;mi++) {
		for (int r=0;r<imgSize;r++) {
			for (int c=0;c<imgSize;c++) {
				images[mi](r,c)=(images[mi](r,c))/maxin;
			}
		}
	}
}

void randomInitialize()
{
	//W^(1)
	double w1Scale=sqrt(inputSize);
	for (int l2=0;l2<hiddenSize;l2++) {
		for (int l1=0;l1<inputSize;l1++) {
			W1(l2,l1)=(double)rand()/((double)RAND_MAX)/w1Scale;
		}
	}	
	//b^(1)
	for (int l2=0;l2<hiddenSize;l2++) {
		W1(l2,inputSize)=(double)rand()/((double)RAND_MAX)/w1Scale;
	}

	//W^(2)
	double w2Scale=sqrt(hiddenSize);
	for (int l3=0;l3<outputSize;l3++) {
		for (int l2=0;l2<hiddenSize;l2++) {
			W2(l3,l2)=(double)rand()/((double)RAND_MAX)/w2Scale;
		}
	}
	//b^(2)
	for (int l3=0;l3<outputSize;l3++) {
		W2(l3,hiddenSize)=(double)rand()/((double)RAND_MAX)/w2Scale;
	}
	//rho
	for (int l2=0;l2<hiddenSize;l2++) {
		rho(l2)=0;
	}
}

void randomSample()
{
	int r0=randInt(imgSize-patchSize);
	int c0=randInt(imgSize-patchSize);
	int im=randInt(numImages);
	for(int r=0;r<patchSize;r++){
		for(int c=0;c<patchSize;c++){
			a1(r*patchSize+c)=images[im](r0+r,c0+c);
		}
	}
}

void blockAlignedSample()
{
	int npatch=imgSize/patchSize;
	int r0=randInt(npatch);
	int c0=randInt(npatch);
	int im=randInt(numImages);
	for (int r=0;r<patchSize;r++) {
		for (int c=0;c<patchSize;c++) {
			a1(r*patchSize+c)=images[im](r0*patchSize+r,c0*patchSize+c);
		}
	}
	a1(inputSize)=1;
}

void normalizeSample()
{
	double sum=0;
	for(int i=0;i<inputSize;i++){
		sum+=a1(i)*a1(i);
	}
	sum=sqrt(sum);
	for(int i=0;i<inputSize;i++){
		a1(i)/=sum;
	}
}


void feedForward2()
{
	for (int l2=0;l2<hiddenSize;l2++) {
		a2(l2)=0;
		for (int l1=0;l1<inputSize+1;l1++) {
			a2(l2)+=W1(l2,l1)*a1(l1);
		}
	}
	a2(hiddenSize)=1;
	for (int l3=0;l3<outputSize;l3++) {
		a3(l3)=0;
		for (int l2=0;l2<hiddenSize+1;l2++) {
			a3(l3)+=W2(l3,l2)*a2(l2);
		}
	}
}

void backPropagation2()
{
	//backpropagation
	for (int l3=0;l3<outputSize;l3++) {
		delta3(l3)=-(a1(l3)-a3(l3));
	}
	for (int l2=0;l2<hiddenSize;l2++) {
		delta2(l2)=0;
		for (int l3=0;l3<outputSize;l3++) {
			delta2(l2)+=delta3(l3)*W2(l3,l2);
		}
	}
	for (int l3=0;l3<outputSize;l3++) {
		for (int l2=0;l2<hiddenSize;l2++) {
			W2(l3,l2)-=alpha*(a2(l2)*delta3(l3)+lambda*W2(l3,l2));
		}
		W2(l3,hiddenSize)-=alpha*delta3(l3);
	}
	for (int l2=0;l2<hiddenSize;l2++) {
		for (int l1=0;l1<inputSize;l1++) {
			W1(l2,l1)-=alpha*(a1(l1)*delta2(l2)+lambda*W1(l2,l1));
		}
		W1(l2,inputSize)-=alpha*delta2(l2);
	}
	//sparsity update
	for (int l2=0;l2<hiddenSize;l2++) {
		rho(l2)=rhoDecay*rho(l2)+(1-rhoDecay)*a2(l2);
	}
// 	for (int l2=0;l2<hiddenSize;l2++) {
// 		W1(l2,inputSize)-=alpha*beta*(rho(l2)-ro);
// 	}
}

void feedForward()
{
	for (int l2=0;l2<hiddenSize;l2++) {
		a2(l2)=0;
		for (int l1=0;l1<inputSize+1;l1++) {
			a2(l2)+=W1(l2,l1)*a1(l1);
		}
		a2(l2)=tanh(a2(l2));
	}
	a2(hiddenSize)=1;
	for (int l3=0;l3<outputSize;l3++) {
		a3(l3)=0;
		for (int l2=0;l2<hiddenSize+1;l2++) {
			a3(l3)+=W2(l3,l2)*a2(l2);
		}
		a3(l3)=tanh(a3(l3));
	}
}

void backPropagation()
{
	//backpropagation
	for (int l3=0;l3<outputSize;l3++) {
		delta3(l3)=-(a1(l3)-a3(l3))*(1.0-a3(l3)*a3(l3));
	}
	for (int l2=0;l2<hiddenSize;l2++) {
		delta2(l2)=0;
		for (int l3=0;l3<outputSize;l3++) {
			delta2(l2)+=delta3(l3)*W2(l3,l2);
		}
		delta2(l2)*=(1.0-a2(l2)*a2(l2));
	}
	for (int l3=0;l3<outputSize;l3++) {
		for (int l2=0;l2<hiddenSize;l2++) {
			W2(l3,l2)-=alpha*(a2(l2)*delta3(l3)+lambda*W2(l3,l2));
		}
		W2(l3,hiddenSize)-=alpha*delta3(l3);
	}
	for (int l2=0;l2<hiddenSize;l2++) {
		for (int l1=0;l1<inputSize;l1++) {
			W1(l2,l1)-=alpha*(a1(l1)*delta2(l2)+lambda*W1(l2,l1));
		}
		W1(l2,inputSize)-=alpha*delta2(l2);
	}
	//sparsity update
	for (int l2=0;l2<hiddenSize;l2++) {
		rho(l2)=rhoDecay*rho(l2)+(1-rhoDecay)*a2(l2);
	}
	for (int l2=0;l2<hiddenSize;l2++) {
		W1(l2,inputSize)-=alpha*beta*(rho(l2)-ro);
	}
}

void printStat()
{
	double err=0;
	cout<<"err: ";
	for (int l3=0;l3<outputSize;l3++) {
		double e=(a3(l3)-a1(l3));
		cout<<e<<' ';
		err+=e*e;
	}
	cout<<endl;
	cout<<"total error: "<<err<<endl;
	cout<<"rho^2: ";
	for (int l2=0;l2<hiddenSize;l2++) {
		cout<<rho(l2)<<' ';
	}
	cout<<endl;
	cout<<"b^1: ";
	for (int l2=0;l2<hiddenSize;l2++) {
		cout<<W1(l2,inputSize)<<' ';
	}
	cout<<endl;
	cout<<"a2: ";
	for (int l2=0;l2<hiddenSize;l2++) {
		//if(a2(l2)>0)
		cout<<a2(l2)<<" ";
	}
	cout<<endl;
}

void outputWeight(int id)
{
	ofstream outdata;	
	sprintf(fstr,"base/weight1_%d.dat",id);
	outdata.open(fstr);
	for (int l2=0;l2<hiddenSize;l2++) {
		for (int l1=0;l1<inputSize;l1++) {
			outdata<<W1(l2,l1)<<" ";
		}
		outdata<<endl;
	}
	outdata.close();
	//	ofstream outdata2;
	sprintf(fstr,"base/weight2_%d.dat",id);
	outdata.open(fstr);
	for (int l2=0;l2<hiddenSize;l2++) {
		for (int l3=0;l3<outputSize;l3++) {
			outdata<<W2(l3,l2)<<" ";
		}
		outdata<<endl;
	}
	outdata.close();
}

void saveData()
{
	ofstream saver;
	char fn[20];
	sprintf(fn,"%s.state",parafile);
	saver.open(fn);
	for(int l2=0;l2<hiddenSize;l2++){
		for(int l1=0;l1<inputSize+1;l1++){
			saver<<W1(l2,l1)<<" ";
		}
		saver<<endl;
	}
	for(int l2=0;l2<hiddenSize+1;l2++){
		for(int l3=0;l3<outputSize;l3++){
			saver<<W2(l3,l2)<<" ";
		}
		saver<<endl;
	}
	for (int l2=0;l2<hiddenSize;l2++) {
		saver<<rho(l2)<<' ';
	}
	saver.close();
}

void loadData()
{
	ifstream loader;
	char fn[20];
	sprintf(fn,"%s.state",parafile);
	loader.open(fn);
	if (!loader) {
		cerr << "Error: file could not be opened" << endl;
		return;
	}
	double num;
	for(int l2=0;l2<hiddenSize;l2++){
		for(int l1=0;l1<inputSize+1;l1++){
			loader>>num;
			W1(l2,l1)=num;
		}
	}
	for(int l2=0;l2<hiddenSize+1;l2++){
		for(int l3=0;l3<outputSize;l3++){
			loader>>num;
			W2(l3,l2)=num;
		}
	}
	for (int l2=0;l2<hiddenSize;l2++) {
		loader>>num;
		rho(l2)=num;
	}
	loader.close();
}

void readParameters()
{
	ifstream loader;
	char fn[20];
	sprintf(fn,"%s.para",parafile);
	loader.open(fn);
	if (!loader) {
		cerr << "Error: file could not be opened" << endl;
		return;
	}
	loader>>loadinit;
	cout<<"loadinit: "<<loadinit<<endl;
	loader>>it;
	cout<<"it: "<<it<<endl;
	loader>>outputnum;
	cout<<"outputnum: "<<outputnum<<endl;
	loader>>alpha;
	cout<<"alpha: "<<alpha<<endl;
	loader>>beta;
	cout<<"beta: "<<beta<<endl;
	loader>>ro;
	cout<<"ro: "<<ro<<endl;
	loader>>lambda;
	cout<<"lambda: "<<lambda<<endl;
	loader>>rhoDecay;
	cout<<"rhoDecay: "<<rhoDecay<<endl;
	loader>>printintv;
	cout<<"printintv: "<<printintv<<endl;
	loader.close();
	
}



int main(int argc, char* argv[])
{

    /**************  Read in images *****************/
	parafile=argv[1];
	cout<<parafile<<endl;
	readParameters();

    
	readImage();
    
    /*************** YOUR CODE HERE ************/
    //normalizeInput();
	//zeroMeanScaleInput();
	//trimInput();
	//tanhNormalize();

	if(loadinit==1)
		loadData();
	else
		randomInitialize();
	
    //start iteration
    for (int i=0;i<it;i++) {
		
		blockAlignedSample();
		//normalizeSample();
        feedForward2();
        backPropagation2();
        
        //for debug
        if (i%printintv==0) {
			cout<<"--------Iteration "<<i<<"-------------"<<endl;
			printStat();
        }
        int intv=it/outputnum;
        if (i%intv==0) {			
            outputWeight(i/intv);
        }
    }
    //save data
	saveData();
    return 0;
}

void activationFunction(const MatrixXd& x, MatrixXd& result) {
    result.resize(x.rows(),x.cols());

    for (int r = 0; r < x.rows(); r++)
        for (int c = 0; c < x.cols(); c++)
            result(r,c) = tanh(x(r,c));
}

int randInt(int lim) {
    return int((double)lim*(double)rand()/(RAND_MAX+1.0));
}
